'''
 _   _                                     _   _                   _          ____               _
| | | |  _   _    __ _   _ __     __ _    | | | |   __ _    ___   ( )___     / ___|   ___     __| |   ___ 
| |_| | | | | |  / _` | | '_ \   / _` |   | |_| |  / _` |  / _ \  |// __|   | |      / _ \   / _` |  / _ \
|  _  | | |_| | | (_| | | | | | | (_| |   |  _  | | (_| | | (_) |   \__ \   | |___  | (_) | | (_| | |  __/
|_| |_|  \__,_|  \__,_| |_| |_|  \__, |   |_| |_|  \__,_|  \___/    |___/    \____|  \___/   \__,_|  \___|
                                  |___/
'''
import numpy as np
import os
from os import system
from astropy.time import Time
import matplotlib.pyplot as plt
import time

from orbfit_pkg._orbfit_running import *

# rwo_file_name应当是 target_origin.rwo
class rwo_file():
    def __init__(self, rwo_file_path: str, rwo_file_name: str, main_work_path = "./"):
        self.rwo_file_path = rwo_file_path
        self.rwo_file_name = rwo_file_name
        self.rwo_file_path_name = os.path.join(self.rwo_file_path, self.rwo_file_name)
        self.object_name = rwo_file_name.split("_")[0]
        self.main_work_path = main_work_path
        # 提取信息
        with open(self.rwo_file_path_name, "r") as rwo_file:
            rwo_file_list = rwo_file.readlines()
            self.rwo_file_list_header = rwo_file_list[:7]
            self.rwo_file_list_data_origin  = rwo_file_list[7:]
            # ========== 对有卫星观测的结果舍去对应行的信息 ==========
            rwo_file_list_data = []
            for info_per_line in self.rwo_file_list_data_origin:
                if len(info_per_line) > 190:
                    rwo_file_list_data.append(info_per_line)
                    
            self.rwo_file_list_data = rwo_file_list_data
            self.changed_rwo_file_list_data = self.rwo_file_list_data
    

    
    '''
    Function :  read_specific_info_col
    Description:
                按照名称和位置提取每一行的信息
                    e.g. read_specific_info_col(rwo_file_list_data,"object", 0, 11)
                返回每一名称信息的所有观测点的字典
                    e.g.{"object", ["469219", "469219", ......]}
    Edition  :  Huang Hao/2023.05.16/
    Log      :
    Input    :  
                info_name               str     information name
                info_index_begin        int     index of begin position per line
                info_index_end_plus1    int     index of end position per line(does not include this position index)
    return   :  
                info_dict               dict    dict of specifit info
    func used:
    '''
    def read_specific_info_col(self, info_name: str, info_index_begin: int, info_index_end_plus1: int) -> dict:
        info_dict = {info_name: "temporary"}
        info_dict_value_list = []
        for per_line_data in self.rwo_file_list_data:
            info_single = per_line_data[info_index_begin: info_index_end_plus1]
            info_dict_value_list.append(info_single)
        info_dict[info_name] = info_dict_value_list
        
        return info_dict
    
    
    '''
    Function :
                change_info_per_line
    Description :
                根据数据在每一行的位置修改数值
    Edition  :  Huang Hao/2023.05.24/
    Log      :
    Input    :
                info_index_begin        int     index of begin position per line
                info_index_end_plus1    int     index of end position per line(does not include this position index)
                info_rwo_index          int     index of which line to change
                new_value               str     what value to change
    Return   :
                None  
    func used:
    Warning  :                
    '''
    def change_info_per_line(self, info_index_begin: int, info_index_end_plus1: int, info_row_index: int, new_value: str) -> None:
        # ========== 提取要修改的一行信息 ==========
        line_to_change = self.changed_rwo_file_list_data[info_row_index]
        # ========== 根据要修改的信息分成前后两端 ==========
        line_first_half = line_to_change[:info_index_begin]
        # print(line_first_half)
        line_last_half = line_to_change[info_index_end_plus1:]
        # print(line_last_half)
        # ========== 替换要修改的信息 ==========
        new_line = line_first_half + new_value + line_last_half
        # ========== 在修改后的rwo文件里生成新的信息 ==========
        self.changed_rwo_file_list_data[info_row_index] = new_line
        
        pass
    
    '''
    Function :  run_orbfit_supplement
    Description:
                使用带supplement的.rwo文件进行定轨计算
    Edition  :  Huang Hao/2023.05.16/
    Log      :
    Input    :  
                back_path               str     How to return to path now
                supplement              str     Supplement for origin rwo file e.g. "_no_2004" >>> 469219_no_2004.rwo
                have_od_oop_file        bool    If already write .oop file >>> TRUE
    return   :  
                info_dict               dict    dict of specifit info
    func used:
    '''   
    def run_orbfit_supplement(self, back_path: str, supplement: str, have_od_oop_file: bool) -> None:

        rwo_supp_file_name = self.object_name+supplement+".rwo"
        
        
        rwo_supp_file_path_name = os.path.join(self.rwo_file_path, rwo_supp_file_name)
        system("cp %s %s.rwo"%(rwo_supp_file_path_name, os.path.join(self.main_work_path, self.object_name)))
        system("rm %s.obs"%(self.object_name))
        # print("cp %s %s.rwo"%(rwo_supp_file_path_name, os.path.join(self.rwo_file_path, self.object_name))); time.sleep(10)
        run_orbfit_supplement(working_path=self.main_work_path, back_path=back_path, object_name=self.object_name,
                                supplement=supplement, have_od_oop_file=have_od_oop_file)
        system("cp %s.rwo %s"%(os.path.join(self.main_work_path, self.object_name), rwo_supp_file_path_name))
        
    '''
    Function :  draw_residual
    Description:
                绘制残差文件
    Edition  :  Huang Hao/2023.05.16/
    Log      :
    Input    :  
    return   :  

    func used:
    '''   
    def draw_residual(self, title_name, figure_path, figure_name, origion_obs_num,own_data_index_list):
        date_dict = self.read_specific_info_col(info_name="Date", info_index_begin=17, info_index_end_plus1=40)
        ra_res_dict = self.read_specific_info_col(info_name="Ra_res", info_index_begin=96, info_index_end_plus1=103)
        dec_res_dict = self.read_specific_info_col(info_name="Dec_res", info_index_begin=149, info_index_end_plus1=156)
        obs_name_dict = self.read_specific_info_col(info_name="Obs_Code", info_index_begin=180, info_index_end_plus1=183)
        dec_dms_dict = self.read_specific_info_col(info_name="Dec_hms", info_index_begin=103, info_index_end_plus1=117)
        
        title = title_name
        date_list = date_dict["Date"]
        for i in range(len(date_list)):
            rwo_date = date_list[i]
            time_isot = convert_rwo_time2isot(rwo_date)
            t = Time(time_isot, format='isot', scale="utc")
            t_decimalyear = t.decimalyear
            date_list[i] = t_decimalyear
        time_year_list = date_list
        
        
        ra_res_list = list(np.float64(ra_res_dict["Ra_res"]))
        dec_res_list = list(np.float64(dec_res_dict["Dec_res"]))
        obs_name_list = obs_name_dict["Obs_Code"]
        dec_dms_list = dec_dms_dict["Dec_hms"]
        # ++++++++++++++++++++++++++++++++++++++++ 将赤纬方向的残差乘cos\delta ++++++++++++++++++++++++++++++++++++++++
        for i in range(len(dec_dms_list)):
            dms = dec_dms_list[i]   #+30 22 01.90  
            degree = int(dms[:3])
            minute = int(dms[4:6])
            second = float(dms[7:])
            dec_deg  =degree + minute/60 + second/3600
            # ra_res_list[i] = ra_res_list[i]*np.cos(np.radians(dec_deg))
            ra_res_list[i] = ra_res_list[i]
            
            
        
        
        size_legend = 10; size_title = 15; size_label = 15; size_tick = 15
        font_legend = {'family' : 'Times New Roman',
                    'weight' : 'heavy',
                    'size'   : size_legend,}
        font_title = {'family' : 'Times New Roman',
                    'weight' : 'heavy',
                    'size'   : size_title,}
        font_label = {'family' : 'Times New Roman',
                    'weight' : 'heavy',
                    'size'   : size_label,}
        font_tick = {'family' : 'Times New Roman',
                        'weight' : 'heavy',
                        'size'   : size_tick,}
        fig, ax = plt.subplots(2, 1, figsize = (6,3), sharex = True, sharey = True)
        plt.subplots_adjust(hspace=0,wspace=0)
        tex_box_props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
        ax[0].scatter(time_year_list, ra_res_list,s=6, label="MPC data",zorder=100)
        
        ax[0].set_ylabel(r'$\Delta\alpha{cos\delta} (^{\prime \prime})$', fontdict=font_label)
        ra_res_mean = np.mean(ra_res_list)
        ra_res_rms = calculate_list_rms(ra_res_list)
        
        ax[0].text(0.05, 0.95, "RMS =%.3f"%(ra_res_rms), transform=ax[0].transAxes, fontdict=font_legend,
                verticalalignment='top', bbox=tex_box_props)
        
        
        ax[0].set_title(title, fontdict=font_title)
        
        ax[1].scatter(time_year_list, dec_res_list,s=6,zorder=100,label="MPC data")
        
        
        if len(time_year_list[origion_obs_num:]) != 0:  # 如果有自主观测数据的话需要把这个画出来
            for index_num in range(len(own_data_index_list)):
                index = own_data_index_list[index_num]
                if index_num == 0:
                    ax[0].scatter(time_year_list[index], ra_res_list[index],s=6, c="r", label = "Own data",zorder=100)
                    ax[1].scatter(time_year_list[index], dec_res_list[index],s=6, c="r", label = "Own data",zorder=100)
                else:
                    ax[0].scatter(time_year_list[index], ra_res_list[index],s=6, c="r", zorder=100)
                    ax[1].scatter(time_year_list[index], dec_res_list[index],s=6, c="r",zorder=100)
        ax[0].legend(prop=font_legend, ncol=2, loc="lower left",frameon=False)
        
        ax[1].set_ylabel(r'$\Delta\delta (^{\prime \prime})$', fontdict=font_label)
        ax[1].set_xlabel('Time / Year', fontdict=font_label)
        dec_res_mean = np.mean(dec_res_list)
        dec_res_rms = calculate_list_rms(dec_res_list)
        ax[1].text(0.05, 0.95, "RMS =%.3f"%(dec_res_rms), transform=ax[1].transAxes, fontdict=font_legend,
                verticalalignment='top', bbox=tex_box_props)
        
        # print("ra res own: ",ra_res_list[origion_obs_num:])
        # print("dec res own: ",dec_res_list[origion_obs_num:])
        
        
        # print("ra: own obs rms:%.2f"%(calculate_list_rms(ra_res_list[origion_obs_num:])))
        # print("dec: own obs rms:%.2f"%(calculate_list_rms(dec_res_list[origion_obs_num:])))
        # print(self.rwo_file_name+"\n", "RMS R.A.=%.5f"%(ra_res_rms), "RMS =%.5f"%(dec_res_rms))
        # ax[0].set_yticks([-0.5, 0, 0.5, 1])
        # ax[1].set_yticks([-1 ,-0.5, 0, 0.5, 1])
        for sub_ax in ax:
            sub_ax.tick_params(axis='both',direction="in",labelsize = size_tick)
            # 设置刻度字体
            x_labels = sub_ax.get_xticklabels()
            y_labels = sub_ax.get_yticklabels()
            sub_ax.grid(zorder=0)
            [label.set_fontname('Times New Roman') for label in x_labels]
            [label.set_fontname('Times New Roman') for label in y_labels]
            [label.set_fontweight('bold') for label in x_labels]
            [label.set_fontweight('bold') for label in y_labels]
        if not os.path.exists(figure_path):
            os.makedirs(figure_path)
        
        plt.savefig(os.path.join(figure_path,figure_name)+".png",bbox_inches="tight", format="png",dpi=1000)
        # plt.savefig(os.path.join(figure_path,figure_name)+".svg",bbox_inches="tight", format="svg",dpi=1000)
        # plt.savefig(os.path.join(figure_path,figure_name)+".eps",bbox_inches="tight", format="eps")
        # plt.savefig(os.path.join(figure_path,figure_name)+".pdf",bbox_inches="tight", format="pdf",dpi=1000)
          

        
'''
Function :  convert_rwo_time2isot
Description:
            将rwo文件中的时间格式转化为isot格式
Edition  :  Huang Hao/2023.05.16/
Log      :
Input    :  
            rwo_time        str     Time in rwo file | e.g."2004 03 17.470006"
return   :  
            time_isot       str     Time in isot format
func used:
    '''   
def convert_rwo_time2isot(rwo_time):
    year = int(rwo_time[:4])
    month = int(rwo_time[5:7])
    day = float(rwo_time[8:])
    
    # 将天的小数部分转换至hms
    day_int = int(day)
    day_float = day - day_int
    # Calculate hour
    hour = day_float * 24
    hour_int = int(hour)
    # Calculate minute
    hour_float = hour - hour_int
    minute = hour_float * 60
    minute_int = int(minute)
    # Calculate second
    minute_float = minute - minute_int
    second = minute_float * 60

    time_isot = "%d-%02d-%02dT%02d:%02d:%02f"%(year, month, day_int, hour_int, minute_int, second)
    
    return time_isot
'''
Function :  write_new_rwo_file
Description:
            写入新的rwo文件
Edition  :  Huang Hao/2023.05.16/
Log      :
Input    :  
            new_rwo_file_dir        str
            new_rwo_file_name       str
            new_rwo_list_header     str
            new_rwo_list_data       str
return   :  

func used:
'''  
def write_new_rwo_file(new_rwo_file_dir: str, new_rwo_file_name: str, new_rwo_list_header: list, new_rwo_list_data: list) -> None:
    new_rwo_file_path_name = os.path.join(new_rwo_file_dir, new_rwo_file_name)
    with open(new_rwo_file_path_name, "w") as new_rwo_file:
        for header_info in new_rwo_list_header:
            new_rwo_file.write(header_info)
        for data_info in new_rwo_list_data:
            new_rwo_file.write(data_info)



def calculate_list_rms(input_list):
    rms = 0
    length = len(input_list)
    for i in range(length):
        rms += input_list[i]**2
        pass
    rms /= length
    rms = np.sqrt(rms)
    
    return rms
if __name__ == "__main__":
    # # ================== Read rwo file and the res value ======================
    # rwo_469219_origion = rwo_file(rwo_file_path="./../test/2016HO3/", rwo_file_name="469219_origin.rwo")
    # # number_dict = rwo_469219.read_specific_info_col(info_name="Number", info_index_begin=0, info_index_end_plus1=11)
    # # date_dict = rwo_469219.read_specific_info_col(info_name="Date", info_index_begin=17, info_index_end_plus1=40)
    # ra_res_dict = rwo_469219_origion.read_specific_info_col(info_name="Ra_res", info_index_begin=96, info_index_end_plus1=103)
    # dec_res_dict = rwo_469219_origion.read_specific_info_col(info_name="Dec_res", info_index_begin=149, info_index_end_plus1=156)

    
    # # =============== Generate the index with specific condition ====================
    # # mean and std
    # ra_res = np.float64(np.array(ra_res_dict["Ra_res"]))
    # ra_res_mean = np.mean(ra_res)
    # ra_res_std = np.std(ra_res)

    # dec_res = np.float64(np.array(dec_res_dict["Dec_res"]))
    # dec_res_mean = np.mean(dec_res)
    # dec_res_std = np.std(dec_res)

    # # 3sigma condition
    # index_3sigma_ra = np.where( ((ra_res < ra_res_mean - 3*ra_res_std) | (ra_res > 3*ra_res_mean + ra_res_std)) )
    # index_3sigma_dec = np.where( ((dec_res < dec_res_mean - 3*dec_res_std) | (dec_res > 3*dec_res_mean + dec_res_std)) )

    # same_index_to_remove_3sigma = np.intersect1d(index_3sigma_ra, index_3sigma_dec)
    # all_index_to_remove_3sigma = np.unique(np.hstack((index_3sigma_ra, index_3sigma_dec)))
    
    # # ====================  Generate and run new rwo file  ==========================
    # # Header of original rwo file
    # rwo_header_list = rwo_469219_origion.rwo_file_list_header
    # # Data of original rwo data
    # rwo_data_list = rwo_469219_origion.rwo_file_list_data
    # rwo_data_list_no_2004 = rwo_data_list[2:]
    
    # # Generate new rwo_data_list according to condition
    # all_index_to_remove_3sigma = list(all_index_to_remove_3sigma)
    # rwo_data_list_no_3sigma = rwo_data_list
    # for index in reversed(all_index_to_remove_3sigma):
    #     rwo_data_list_no_3sigma.pop(index)
    
    
    # # Write and run new rwo file: 3sigma reject
    # write_new_rwo_file(new_rwo_file_dir="./../test/2016HO3/", new_rwo_file_name="469219_no_3sigma.rwo",
    #                    new_rwo_list_header=rwo_header_list, new_rwo_list_data=rwo_data_list_no_3sigma)
    # rwo_469219_origion.run_orbfit_supplement(back_path="./..", supplement="_no_3sigma", have_od_oop_file=True)
    
    # # Write and run new rwo file: no_2004
    # write_new_rwo_file(new_rwo_file_dir="./../test/2016HO3/", new_rwo_file_name="469219_no_2004.rwo",
    #                    new_rwo_list_header=rwo_header_list, new_rwo_list_data=rwo_data_list_no_2004)
    # rwo_469219_origion.run_orbfit_supplement(back_path="./../../src", supplement="_no_2004", have_od_oop_file=True)
    
    # #Run new rwo file: neodys
    # rwo_469219_origion.run_orbfit_supplement(back_path="./../../src", supplement="_neodys", have_od_oop_file=True)
    # #Run new rwo file: neodys
    # rwo_469219_origion.run_orbfit_supplement(back_path="./../../src", supplement="_origion", have_od_oop_file=True)
    
    # # ============================  Draw the Res  ===================================
    # rwo_469219_origion = rwo_file(rwo_file_path="./../test/2016HO3/", rwo_file_name="469219_origin.rwo")
    # rwo_469219_origion.draw_residual(title_name=rwo_469219_origion.rwo_file_name.split(".")[0],
    #                                  figure_path=os.path.join(rwo_469219_origion.rwo_file_path, "fig_test/"),
    #                                  figure_name = rwo_469219_origion.rwo_file_name.split(".")[0])

    # rwo_469219_no_2004 = rwo_file(rwo_file_path="./../test/2016HO3/", rwo_file_name="469219_no_2004.rwo")
    # rwo_469219_no_2004.draw_residual(title_name=rwo_469219_no_2004.rwo_file_name.split(".")[0],
    #                                  figure_path=os.path.join(rwo_469219_no_2004.rwo_file_path, "fig_test/"),
    #                                  figure_name = rwo_469219_no_2004.rwo_file_name.split(".")[0])
    
    # rwo_469219_no_3sigma = rwo_file(rwo_file_path="./../test/2016HO3/", rwo_file_name="469219_no_3sigma.rwo")
    # rwo_469219_no_3sigma.draw_residual(title_name=rwo_469219_no_3sigma.rwo_file_name.split(".")[0],
    #                                  figure_path=os.path.join(rwo_469219_no_3sigma.rwo_file_path, "fig_test/"),
    #                                  figure_name = rwo_469219_no_3sigma.rwo_file_name.split(".")[0])
    
    # rwo_469219_neodys = rwo_file(rwo_file_path="./../test/2016HO3/", rwo_file_name="469219_neodys.rwo")
    # rwo_469219_neodys.draw_residual(title_name=rwo_469219_neodys.rwo_file_name.split(".")[0],
    #                                  figure_path=os.path.join(rwo_469219_neodys.rwo_file_path, "fig_test/"),
    #                                  figure_name = rwo_469219_neodys.rwo_file_name.split(".")[0])
    
    
    
    rwo_469219_neodys = rwo_file(rwo_file_path="./2016HO3/", rwo_file_name="469219.rwo")
    rwo_469219_neodys.draw_residual(title_name="TEST_SRP_A2",
                                     figure_path="./",
                                     figure_name = "test_obs_res")

